import torch
from torch import nn

def m(x):
    return -1 + torch.log(1 + torch.exp(x))

def h(x):
    return torch.tanh(x)

def h_prime(x):
    return 1 - torch.tanh(x) ** 2

def plainer(z, ws, bs, logdet=False):
    wz_b = ws[0] * z + bs[0]
    LDJ = -torch.log(torch.abs(ws[0]))
    for w, b in zip(ws[1:], bs[1:]):
        x = h(wz_b)
        LDJ += -torch.log(torch.abs(u * w * h_prime(wz_b)) + 1e-8) # batch x 1
    return x, LDJ

class PlainerFlow:
    def __init__(self, p):
        self.p = p

    def log_prob(x):


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.lstm = nn.LSTM(5, 64, 2, batch_first=True)
        self.out = nn.Linear(64, 3 * 4)

        self.init_weights()

    def init_weights(self):
        nn.init.normal_(self.out.weight, 0.01)
        nn.init.constant_(self.out.bias, 0)

    def forward(self, x, h=None, ret_h=False):
        x, h = self.lstm(x, h)
        x = self.out(x).view(-1, 4, 3)
        d = PlainerFlow(x)
        if ret_h:
            return (d, x[..., -1].sigmoid()), h
        return (d, x[..., -1].sigmoid())
